package com.lzh.crypt.ase;

import com.lzh.crypt.enums.AesAlgorithmEnum;
import com.lzh.crypt.exception.CryptException;
import com.lzh.crypt.exception.ErrorCode;
import org.bouncycastle.jce.provider.BouncyCastleProvider;

import javax.crypto.Cipher;
import javax.crypto.KeyGenerator;
import javax.crypto.SecretKey;
import javax.crypto.spec.IvParameterSpec;
import javax.crypto.spec.SecretKeySpec;
import java.security.*;
import java.security.spec.AlgorithmParameterSpec;

/**
 * @author zhehen.lu
 * @date 2025/8/10 19:55
 */
public class AESAdvancedKeyPair {

    private final byte[] encryptKey;

    private final AlgorithmParameterSpec algorithmParameterSpec;

    private final int mode;

    private final String cryptAlgorithm;

    private final byte[] additionalAuthData;

    static {
        Security.addProvider(new BouncyCastleProvider());
    }

    protected AESAdvancedKeyPair(AesAlgorithmEnum aesAlgorithmEnum, byte[] encryptKey, byte[] ivByte, byte[] additionalAuthData,
                                 int mode) {
        this.encryptKey = encryptKey;
        algorithmParameterSpec = ivByte == null ? null : aesAlgorithmEnum.generrateParameterSpec(128, ivByte);
        this.mode = mode;
        cryptAlgorithm = aesAlgorithmEnum.getAlgorithm();
        this.additionalAuthData = additionalAuthData;
    }

    public static AESAdvancedKeyPair instance(AesAlgorithmEnum aesAlgorithmEnum, byte[] encryptKey, byte[] ivByte, int mode) {
        try {
            return new AESAdvancedKeyPair(aesAlgorithmEnum, encryptKey, ivByte, null, mode);
        } catch (Exception e) {
            throw new CryptException(ErrorCode.ERROR_GENERATING_AES_KEY.getErrorMsg(), e);
        }
    }

    public static AESAdvancedKeyPair instance(AesAlgorithmEnum aesAlgorithmEnum, byte[] encryptKey, byte[] ivByte,
                                              byte[] additionalAuthData, int mode) {
        try {
            return new AESAdvancedKeyPair(aesAlgorithmEnum, encryptKey, ivByte, additionalAuthData, mode);
        } catch (Exception e) {
            throw new CryptException(ErrorCode.ERROR_GENERATING_AES_KEY.getErrorMsg(), e);
        }
    }

    public static AESAdvancedKeyPair instance(AesAlgorithmEnum aesAlgorithmEnum, byte[] encryptKey, int mode) {
        try {
            return new AESAdvancedKeyPair(aesAlgorithmEnum, encryptKey, null, null, mode);
        } catch (Exception e) {
            throw new CryptException(ErrorCode.ERROR_GENERATING_AES_KEY.getErrorMsg(), e);
        }
    }

    public static byte[] genericKey() throws Exception {
        KeyGenerator generator = KeyGenerator.getInstance(AesAlgorithmEnum.AES.getAlgorithm());
        generator.init(128, SecureRandom.getInstanceStrong());
        SecretKey secretKey = generator.generateKey();
        return secretKey.getEncoded();
    }

    public static byte[] genericKey(int size) throws Exception {
        KeyGenerator generator = KeyGenerator.getInstance(AesAlgorithmEnum.AES.getAlgorithm());
        generator.init(size, SecureRandom.getInstanceStrong());
        SecretKey secretKey = generator.generateKey();
        return secretKey.getEncoded();
    }

    public byte[] encrypt(byte[] originalContent) {
        try {
            Cipher cipher = Cipher.getInstance(cryptAlgorithm);
            SecretKeySpec skeySpec = new SecretKeySpec(encryptKey, AesAlgorithmEnum.AES.getAlgorithm());
            cipher.init(mode, skeySpec, algorithmParameterSpec);
            if (additionalAuthData != null) {
                cipher.updateAAD(additionalAuthData);
            }
            byte[] encrypted = cipher.doFinal(originalContent);
            return encrypted;
        } catch (Exception e) {
            throw new CryptException(ErrorCode.ENCRYPTION_ERROR.getErrorMsg(), e);
        }
    }

    /**
     * AES解密 填充模式AES/CBC/PKCS7Padding 解密模式128
     *
     * @param content 目标密文
     * @return {@link byte[]}
     * @date 2025/8/10 14:04
     * @author zhehen.lu
     */
    public byte[] decrypt(byte[] content) {
        try {
            Cipher cipher = Cipher.getInstance(cryptAlgorithm);
            SecretKeySpec skeySpec = new SecretKeySpec(encryptKey, AesAlgorithmEnum.AES.getAlgorithm());
            cipher.init(mode, skeySpec, algorithmParameterSpec);
            if (additionalAuthData != null) {
                cipher.updateAAD(additionalAuthData);
            }
            byte[] result = cipher.doFinal(content);
            return result;
        } catch (Exception e) {
            throw new CryptException(ErrorCode.DECRYPTION_ERROR.getErrorMsg(), e);
        }
    }

    /**
     * 生成iv
     *
     * @param iv
     * @return {@link AlgorithmParameters}
     * @date 2025/8/10 14:04
     * @author zhehen.lu
     */
    public static AlgorithmParameters generateIV(byte[] iv) throws Exception {
        AlgorithmParameters params = AlgorithmParameters.getInstance(AesAlgorithmEnum.AES.getAlgorithm());
        params.init(new IvParameterSpec(iv));
        return params;
    }

}
